{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Automatic Speech Recognition with Speaker Diarization"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "\"\"\"\n",
                "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
                "\n",
                "Instructions for setting up Colab are as follows:\n",
                "1. Open a new Python 3 notebook.\n",
                "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
                "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
                "4. Run this cell to set up dependencies.\n",
                "\"\"\"\n",
                "# If you're using Google Colab and not running locally, run this cell.\n",
                "\n",
                "## Install dependencies\n",
                "!pip install wget\n",
                "!apt-get install sox libsndfile1 ffmpeg\n",
                "!pip install unidecode\n",
                "\n",
                "# ## Install NeMo\n",
                "BRANCH = 'main'\n",
                "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[asr]\n",
                "\n",
                "## Install TorchAudio\n",
                "!pip install torchaudio -f https://download.pytorch.org/whl/torch_stable.html"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Introduction"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Speaker diarization lets us figure out \"who spoke when\" in the transcription. Without speaker diarization, we cannot distinguish the speakers in the transcript generated from automatic speech recognition (ASR). Nowadays, ASR combined with speaker diarization has shown immense use in many tasks, ranging from analyzing meeting transcription to media indexing. \n",
                "\n",
                "In this tutorial, we demonstrate how we can get ASR transcriptions combined with speaker labels. Since we don't include a detailed process of getting ASR results or diarization results, please refer to the following links for more in-depth description.\n",
                "\n",
                "If you need detailed understanding of transcribing words with ASR, refer to this [ASR Tutorial](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb) tutorial.\n",
                "\n",
                "\n",
                "For detailed parameter setting and execution of speaker diarization, refer to this [Diarization Inference](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/speaker_tasks/Speaker_Diarization_Inference.ipynb) tutorial.\n",
                "\n",
                "\n",
                "An example script that runs ASR and speaker diarization together can be found at [ASR with Diarization](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/diarization/offline_diarization_with_asr.py).\n",
                "\n",
                "### Speaker diarization in ASR pipeline\n",
                "\n",
                "Speaker diarization results in ASR pipeline should align well with ASR output. Thus, we use ASR output to create Voice Activity Detection (VAD) timestamps to obtain segments we want to diarize. The segments we obtain from the VAD timestamps are further segmented into sub-segments in the speaker diarization step. Finally, after obtaining the speaker labels from speaker diarization, we match the decoded words with speaker labels to generate a transcript with speaker labels.\n",
                "\n",
                "    ASR → VAD timestamps and decoded words → speaker diarization → speaker label matching"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Import libraries\n",
                "\n",
                "Let's first import nemo asr and other libraries for visualization purposes."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import nemo.collections.asr as nemo_asr\n",
                "import numpy as np\n",
                "from IPython.display import Audio, display\n",
                "import librosa\n",
                "import os\n",
                "import wget\n",
                "import matplotlib.pyplot as plt\n",
                "\n",
                "import nemo\n",
                "import glob\n",
                "\n",
                "import pprint\n",
                "pp = pprint.PrettyPrinter(indent=4)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "We demonstrate this tutorial using a merged AN4 audioclip. The merged audioclip contains the speech of two speakers (male and female) reading dates in different formats. Run the following script to download the audioclip and play it."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "ROOT = os.getcwd()\n",
                "data_dir = os.path.join(ROOT,'data')\n",
                "os.makedirs(data_dir, exist_ok=True)\n",
                "\n",
                "an4_audio_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.wav\"\n",
                "if not os.path.exists(os.path.join(data_dir,'an4_diarize_test.wav')):\n",
                "    AUDIO_FILENAME = wget.download(an4_audio_url, data_dir)\n",
                "else:\n",
                "    AUDIO_FILENAME = os.path.join(data_dir,'an4_diarize_test.wav')\n",
                "\n",
                "audio_file_list = glob.glob(f\"{data_dir}/*.wav\")\n",
                "print(\"Input audio file list: \\n\", audio_file_list)\n",
                "\n",
                "signal, sample_rate = librosa.load(AUDIO_FILENAME, sr=None)\n",
                "display(Audio(signal,rate=sample_rate))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "`display_waveform()` and `get_color()` functions are defined for displaying the waveform with diarization results."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def display_waveform(signal,text='Audio',overlay_color=[]):\n",
                "    fig,ax = plt.subplots(1,1)\n",
                "    fig.set_figwidth(20)\n",
                "    fig.set_figheight(2)\n",
                "    plt.scatter(np.arange(len(signal)),signal,s=1,marker='o',c='k')\n",
                "    if len(overlay_color):\n",
                "        plt.scatter(np.arange(len(signal)),signal,s=1,marker='o',c=overlay_color)\n",
                "    fig.suptitle(text, fontsize=16)\n",
                "    plt.xlabel('time (secs)', fontsize=18)\n",
                "    plt.ylabel('signal strength', fontsize=14);\n",
                "    plt.axis([0,len(signal),-0.5,+0.5])\n",
                "    time_axis,_ = plt.xticks();\n",
                "    plt.xticks(time_axis[:-1],time_axis[:-1]/sample_rate);\n",
                "    \n",
                "COLORS=\"b g c m y\".split()\n",
                "\n",
                "def get_color(signal,speech_labels,sample_rate=16000):\n",
                "    c=np.array(['k']*len(signal))\n",
                "    for time_stamp in speech_labels:\n",
                "        start,end,label=time_stamp.split()\n",
                "        start,end = int(float(start)*16000),int(float(end)*16000),\n",
                "        if label == \"speech\":\n",
                "            code = 'red'\n",
                "        else:\n",
                "            code = COLORS[int(label.split('_')[-1])]\n",
                "        c[start:end]=code\n",
                "    \n",
                "    return c "
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Using the above function, we can display the waveform of the example audio clip. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "display_waveform(signal)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Parameter setting for ASR and diarization\n",
                "First, we need to setup the following parameters for ASR and diarization. We start our demonstration by first transcribing the audio recording using our pretrained ASR model `QuartzNet15x5Base-En` and use the CTC output probabilities to get timestamps for the spoken words. We then use these timestamps to get speaker label information using the speaker diarizer model. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from omegaconf import OmegaConf\n",
                "import shutil\n",
                "CONFIG_URL = \"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/offline_diarization_with_asr.yaml\"\n",
                "\n",
                "if not os.path.exists(os.path.join(data_dir,'offline_diarization_with_asr.yaml')):\n",
                "    CONFIG = wget.download(CONFIG_URL, data_dir)\n",
                "else:\n",
                "    CONFIG = os.path.join(data_dir,'offline_diarization_with_asr.yaml')\n",
                "\n",
                "cfg = OmegaConf.load(CONFIG)\n",
                "print(OmegaConf.to_yaml(cfg))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Speaker Diarization scripts commonly expects following arguments:\n",
                "1. manifest_filepath : Path to manifest file containing json lines of format: `{\"audio_filepath\": \"/path/to/audio_file\", \"offset\": 0, \"duration\": null, \"label\": \"infer\", \"text\": \"-\", \"num_speakers\": null, \"rttm_filepath\": \"/path/to/rttm/file\", \"uem_filepath\"=\"/path/to/uem/filepath\"}`\n",
                "2. out_dir : directory where outputs and intermediate files are stored. \n",
                "3. oracle_vad: If this is true then we extract speech activity labels from rttm files, if False then either \n",
                "4. vad.model_path or external_manifestpath containing speech activity labels has to be passed. \n",
                "\n",
                "Mandatory fields are `audio_filepath`, `offset`, `duration`, `label` and `text`. For the rest if you would like to evaluate with a known number of speakers pass the value else `null`. If you would like to score the system with known rttms then that should be passed as well, else `null`. uem file is used to score only part of your audio for evaluation purposes, hence pass if you would like to evaluate on it else `null`.\n",
                "\n",
                "\n",
                "**Note:** we expect audio and corresponding RTTM to have **same base name** and the name should be **unique**. \n",
                "\n",
                "For example: if audio file name is **test_an4**.wav, if provided we expect corresponding rttm file name to be **test_an4**.rttm (note the matching **test_an4** base name)\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Lets create a manifest file with the an4 audio and rttm available. If you have more than one file you may also use the script `NeMo/scripts/speaker_tasks/pathsfiles_to_manifest.py` to generate a manifest file from a list of audio files. In addition, you can optionally include rttm files to evaluate the diarization results."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Create a manifest file for input with below format. \n",
                "# {\"audio_filepath\": \"/path/to/audio_file\", \"offset\": 0, \"duration\": null, \"label\": \"infer\", \"text\": \"-\", \n",
                "# \"num_speakers\": null, \"rttm_filepath\": \"/path/to/rttm/file\", \"uem_filepath\"=\"/path/to/uem/filepath\"}\n",
                "import json\n",
                "meta = {\n",
                "    'audio_filepath': AUDIO_FILENAME, \n",
                "    'offset': 0, \n",
                "    'duration':None, \n",
                "    'label': 'infer', \n",
                "    'text': '-', \n",
                "    'num_speakers': 2, \n",
                "    'rttm_filepath': None, \n",
                "    'uem_filepath' : None\n",
                "}\n",
                "with open(os.path.join(data_dir,'input_manifest.json'),'w') as fp:\n",
                "    json.dump(meta,fp)\n",
                "    fp.write('\\n')\n",
                "\n",
                "cfg.diarizer.manifest_filepath = os.path.join(data_dir,'input_manifest.json')\n",
                "!cat {cfg.diarizer.manifest_filepath}"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Let's set the parameters required for diarization. In this tutorial, we obtain voice activity labels from ASR, which is set through parameter `cfg.diarizer.asr.parameters.asr_based_vad`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "pretrained_speaker_model='titanet_large'\n",
                "cfg.diarizer.manifest_filepath = cfg.diarizer.manifest_filepath\n",
                "cfg.diarizer.out_dir = data_dir #Directory to store intermediate files and prediction outputs\n",
                "cfg.diarizer.speaker_embeddings.model_path = pretrained_speaker_model\n",
                "cfg.diarizer.speaker_embeddings.parameters.window_length_in_sec = 1.5\n",
                "cfg.diarizer.speaker_embeddings.parameters.shift_length_in_sec = 0.75\n",
                "cfg.diarizer.clustering.parameters.oracle_num_speakers=True\n",
                "\n",
                "# Using VAD generated from ASR timestamps\n",
                "cfg.diarizer.asr.model_path = 'QuartzNet15x5Base-En'\n",
                "cfg.diarizer.oracle_vad = False # ----> Not using oracle VAD \n",
                "cfg.diarizer.asr.parameters.asr_based_vad = True\n",
                "cfg.diarizer.asr.parameters.threshold=100 # ASR based VAD threshold: If 100, all silences under 1 sec are ignored.\n",
                "cfg.diarizer.asr.parameters.decoder_delay_in_sec=0.2 # Decoder delay is compensated for 0.2 sec"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Run ASR and get word timestamps\n",
                "Before we run speaker diarization, we should run ASR and get the ASR output to generate decoded words and timestamps for those words. Let's import `ASR_TIMESTAMPS` class and create `asr_ts_decoder` instance that returns an ASR model. Using this ASR model, the following two variables are obtained from `asr_ts_decoder.run_ASR()` function."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "- word_hyp Dict[str, List[str]]: contains the sequence of words.\n",
                "- word_ts_hyp Dict[str, List[int]]: contains frame level index of the start and the end of each word."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from nemo.collections.asr.parts.utils.decoder_timestamps_utils import ASR_TIMESTAMPS\n",
                "asr_ts_decoder = ASR_TIMESTAMPS(**cfg.diarizer)\n",
                "asr_model = asr_ts_decoder.set_asr_model()\n",
                "word_hyp, word_ts_hyp = asr_ts_decoder.run_ASR(asr_model)\n",
                "\n",
                "print(\"Decoded word output dictionary: \\n\", word_hyp['an4_diarize_test'])\n",
                "print(\"Word-level timestamps dictionary: \\n\", word_ts_hyp['an4_diarize_test'])"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Let's create an instance `asr_diar_offline` from ASR_DIAR_OFFLINE class, which matches diarization results with ASR outputs. We pass ``cfg.diarizer`` to setup the parameters for both ASR and diarization. We also set `word_ts_anchor_offset` variable that determines the anchor position of each word. Here, we use the default value from `asr_ts_decoder` instance."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from nemo.collections.asr.parts.utils.diarization_utils import ASR_DIAR_OFFLINE\n",
                "asr_diar_offline = ASR_DIAR_OFFLINE(**cfg.diarizer)\n",
                "asr_diar_offline.word_ts_anchor_offset = asr_ts_decoder.word_ts_anchor_offset"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "`asr_diar_offline` instance is now ready. As a next step, we run diarization."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Run diarization with the extracted word timestamps"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now that all the components for diarization is ready, let's run diarization by calling `run_diarization()` function. `run_diarization()` will return two different variables : `diar_hyp` and `diar_score`. `diar_hyp` is diarization inference result which is written in `[start time] [end time] [speaker]` format. `diar_score` contains `None` since we did not provide `rttm_filepath` in the input manifest file."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "diar_hyp, diar_score = asr_diar_offline.run_diarization(cfg, word_ts_hyp)\n",
                "print(\"Diarization hypothesis output: \\n\", diar_hyp['an4_diarize_test'])"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "`run_diarization()` function also creates `an4_diarize_test.rttm` file. Let's check what is written in this `rttm` file."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def read_file(path_to_file):\n",
                "    with open(path_to_file) as f:\n",
                "        contents = f.read().splitlines()\n",
                "    return contents\n",
                "\n",
                "predicted_speaker_label_rttm_path = f\"{data_dir}/pred_rttms/an4_diarize_test.rttm\"\n",
                "pred_rttm = read_file(predicted_speaker_label_rttm_path)\n",
                "\n",
                "pp.pprint(pred_rttm)\n",
                "\n",
                "from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels\n",
                "pred_labels = rttm_to_labels(predicted_speaker_label_rttm_path)\n",
                "\n",
                "color = get_color(signal, pred_labels)\n",
                "display_waveform(signal,'Audio with Speaker Labels', color)\n",
                "display(Audio(signal,rate=16000))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Check the speaker-labeled ASR transcription output\n",
                "\n",
                "Now we've done all the processes for running ASR and diarization, let's match the diarization result with the ASR result and get the final output. `get_transcript_with_speaker_labels()` function in `asr_diar_offline` matches diarization output `diar_hyp` with `word_hyp` using the timestamp information from `word_ts_hyp`. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "asr_diar_offline.get_transcript_with_speaker_labels(diar_hyp, word_hyp, word_ts_hyp)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "After running `get_transcript_with_speaker_labels()` function, the transcription output will be located in `./pred_rttms` folder, which shows **start time to end time of the utterance, speaker ID, and words spoken** during the notified time."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "transcription_path_to_file = f\"{data_dir}/pred_rttms/an4_diarize_test.txt\"\n",
                "transcript = read_file(transcription_path_to_file)\n",
                "pp.pprint(transcript)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Another output is transcription output in JSON format, which is saved in `./pred_rttms/an4_diarize_test.json`. \n",
                "\n",
                "In the JSON format output, we include information such as **transcription, estimated number of speakers (variable named `speaker_count`), start and end time of each word and most importantly, speaker label for each word.**"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "transcription_path_to_file = f\"{data_dir}/pred_rttms/an4_diarize_test.json\"\n",
                "json_contents = read_file(transcription_path_to_file)\n",
                "pp.pprint(json_contents)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Optional Features for ASR with Speaker Diarization\n",
                "\n",
                "#### Beam search decoder\n",
                "Beam-search decoder can be applied to CTC based ASR models. To use this feature, [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode) should be installed. [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode) supports word timestamp generation and can be applied to speaker diarization. pyctcdecode also requires [KenLM](https://github.com/kpu/kenlm) and KenLM is recommended to be installed using PyPI. Install pyctcdecode in your environment with the following commands:\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "!pip install pyctcdecode\n",
                "!pip install https://github.com/kpu/kenlm/archive/master.zip"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "You can download publicly available language models (`.arpa` files) at [KALDI Tedlium Language Models](https://kaldi-asr.org/models/m5). Download [4-gram Big ARPA](https://kaldi-asr.org/models/5/4gram_big.arpa.gz) and provide the model path. Let's download the language model file to `data_dir` folder."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import gzip\n",
                "import shutil\n",
                "def gunzip(file_path,output_path):\n",
                "    with gzip.open(file_path,\"rb\") as f_in, open(output_path,\"wb\") as f_out:\n",
                "        shutil.copyfileobj(f_in, f_out)\n",
                "        f_in.close()\n",
                "        f_out.close()\n",
                "        \n",
                "ARPA_URL = 'https://kaldi-asr.org/models/5/4gram_big.arpa.gz'\n",
                "f = wget.download(ARPA_URL, data_dir)\n",
                "gunzip(f,f.replace(\".gz\",\"\"))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Provide the downloaded arpa language model file to `cfg.diarizer`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "arpa_model_path = os.path.join(data_dir, '4gram_big.arpa')\n",
                "cfg.diarizer.asr.ctc_decoder_parameters.pretrained_language_model = arpa_model_path"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "create a new `asr_ts_decoder` instance with the updated `cfg.diarizer`. The decoder script will launch pyctcdecode for decoding words and timestamps."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import importlib\n",
                "import nemo.collections.asr.parts.utils.decoder_timestamps_utils as decoder_timestamps_utils\n",
                "importlib.reload(decoder_timestamps_utils) # This module should be reloaded after you install pyctcdecode.\n",
                "\n",
                "asr_ts_decoder = ASR_TIMESTAMPS(**cfg.diarizer)\n",
                "asr_model = asr_ts_decoder.set_asr_model()\n",
                "word_hyp, word_ts_hyp = asr_ts_decoder.run_ASR(asr_model)\n",
                "\n",
                "print(\"Decoded word output dictionary: \\n\", word_hyp['an4_diarize_test'])\n",
                "print(\"Word-level timestamps dictionary: \\n\", word_ts_hyp['an4_diarize_test'])"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Realign Words with a Language Model (Experimental)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Diarization result with ASR transcript can be enhanced by applying a language model. The mapping between speaker labels and words can be realigned by employing language models. The realigning process calculates the probability of the words around the boundary between two hypothetical sentences spoken by different speakers.\n",
                "\n",
                " <Example> k-th word: `but`\n",
                "    \n",
                "            hyp_former:\n",
                "                \"since i think like tuesday </s> <s> but he's coming back to albuquerque\"\n",
                "    \n",
                "            hyp_latter:\n",
                "                \"since i think like tuesday but </s> <s> he's coming back to albuquerque\"\n",
                "\n",
                "The joint probabilities of words in the sentence are computed for these two hypotheses. In this example, `hyp_former` is likely to get a higher score and thus word `but` will be assigned to the second speaker.\n",
                "\n",
                "To use this feature, python package [arpa](https://pypi.org/project/arpa/) should be installed."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "!pip install arpa"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "`diarizer.asr.realigning_lm_parameters.logprob_diff_threshold` can be modified to optimize the diarization performance (default value is 1.2). This is a threshold value for the gap between two log-probabilities of two hypotheses. Thus, the lower the threshold, the more changes are expected to be seen in the output transcript.   \n",
                "\n",
                "`arpa` package also uses KenLM language models as in pyctcdecode. You can download publicly available [4-gram Big ARPA](https://kaldi-asr.org/models/5/4gram_big.arpa.gz) model and provide the model path to hydra configuration as follows.\n",
                " "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "arpa_model_path = os.path.join(data_dir, '4gram_big.arpa')\n",
                "cfg.diarizer.asr.realigning_lm_parameters.arpa_language_model = arpa_model_path\n",
                "cfg.diarizer.asr.realigning_lm_parameters.logprob_diff_threshold = 1.2\n",
                "\n",
                "import importlib\n",
                "import nemo.collections.asr.parts.utils.diarization_utils as diarization_utils\n",
                "importlib.reload(diarization_utils) # This module should be reloaded after you install arpa.\n",
                "\n",
                "# Create a new instance with realigning language model\n",
                "asr_diar_offline = ASR_DIAR_OFFLINE(**cfg.diarizer)\n",
                "asr_diar_offline.word_ts_anchor_offset = asr_ts_decoder.word_ts_anchor_offset"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now that the language model for realigning is set up, you can run `get_transcript_with_speaker_labels()` to get the results with realigning. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "asr_diar_offline.get_transcript_with_speaker_labels(diar_hyp, word_hyp, word_ts_hyp)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "transcription_path_to_file = f\"{data_dir}/pred_rttms/an4_diarize_test.txt\"\n",
                "transcript = read_file(transcription_path_to_file)\n",
                "pp.pprint(transcript)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.8.8"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}